
import sys,os,os.path,getopt

class FO_H:
	def write(self,S):
		pass
	def close(self):
		pass

def Main(argv):
	(options,args) = getopt.getopt(argv[1:],"Vsm:?")
	modname = os.path.splitext(args[0])[0]

	for i,j in options:
		if i == 'm':
			modname = j

	if len(args) > 1 :
		asmfile = args[1]
	else:	
		asmfile = modname + '.S'

	f = open(args[0],"r")
	fo_S = open(asmfile,"w+")
	modname = os.path.basename(modname)
	#sys.stderr.write(asmfile+" "+modname+"\n")
	if modname == "ntdll":
		fo_H = open(modname+'.H',"w+")
	else:
		fo_H = FO_H() #open(modname+'.H',"w+")
	if len(args) > 2:
		section = args[2]
	else:
		section = '.text$i'
	fn_no = 0
	fn_names = []

	fo_H.write(\
'''
/*

  This file autogenerated by gen_autoimport_table_py from '''+args[0]+'''
  Copyright (C) 2007-2009, Alexey Sudachen, Teggo Software Ltd.

*/

#ifndef ___85c4fbf2_8e26_4b8f_8670_524925929cd4___'''+modname+'''
#define ___85c4fbf2_8e26_4b8f_8670_524925929cd4___'''+modname+'''

#ifdef __cplusplus
extern "C" {
#endif

extern HANDLE _auto_iat_'''+modname+'''_handle;
extern void '''+("_%s_xcall_fullinit"%modname)+'''();
''');

	fo_S.write(\
'''
/*

  This file autogenerated by gen_autoimport_table_py from '''+args[0]+'''
  Copyright (C) 2007-2009, Alexey Sudachen, alexey@sudachen.name.

*/

.text
/*.section '''+section+'''*/

''')

	imp_names = {}
	imp2_names = {}
	xcd_names = {}

	for i in f.xreadlines():
		i = i.strip()
		if not i or i[0] == '#' : continue

		i = i.split(',')
		funame = i[0]
		xcdecl = False
		if funame.startswith('%'):
			xcdecl = True
			funame = funame[1:]
		rtype = i[1]
		arg_types = i[2:]
		if funame.startswith('-'):
			funame = funame[1:]
		if not xcdecl:
			imp_names[fn_no] = '%s@%d' % (funame,len(arg_types)*4)
			imp2_names[fn_no] = '%s_%d' % (funame,len(arg_types)*4)
			fo_H.write("%s __stdcall _%s_%s_xcall_%d(int _xcall_no" % (rtype,modname,funame,fn_no))
			j = 0
			for i in arg_types:
				fo_H.write(",%s _xcall_arg%d" % (i,j))
				j += 1
			fo_H.write(");\n")
			fo_H.write("#undef %s\n" % funame)
			fo_H.write("#define %s(" % funame)
			j = ''
			for i in xrange(len(arg_types)):
				fo_H.write("%s_xcall_x%d" % (j,i))
				j = ','
			fo_H.write(") \\\n  _%s_%s_xcall_%d(%d\\\n" % (modname,funame,fn_no,fn_no))
			fo_H.write("    ")
			for i in xrange(len(arg_types)):
				fo_H.write(",_xcall_x%d" % i)
			fo_H.write(")\n")

			fo_S.write(".global __%s_%s_xcall_%d\n" %(modname,funame,fn_no))
			fo_S.write("__%s_%s_xcall_%d:\n" %(modname,funame,fn_no))
			fo_S.write(".global __%s_%s_xcall_%d@%d\n" %(modname,funame,fn_no,len(arg_types)*4+4))
			fo_S.write("__%s_%s_xcall_%d@%d:\n" %(modname,funame,fn_no,len(arg_types)*4+4))
		else:
			xcd_names[fn_no] = funame

		fn_names.append(funame)

		fn_no += 1

	fo_H.write(\
'''
#ifdef __cplusplus
} /* extern C */
#endif

#endif /*___85c4fbf2_8e26_4b8f_8670_524925929cd4___'''+modname+'''*/

''');

	fo_H.close()

	fo_S.write(\
'''
	pop     %eax
	xchg    (%esp),%eax
'''+("__%s_xcall_resolver_X"%modname)+''':
	shll    $2,%eax
	lea     '''+("__%s_TABLE"%modname)+'''(%eax),%eax
	jmp     *(%eax)

'''+("__%s_xcall_resolver"%modname)+''':
	push    %eax
	sub     $'''+("__%s_TABLE"%modname)+''',%eax
	push    %eax
	lea     '''+("__%s_TABLE_1"%modname)+'''(%eax),%eax
	mov     (%eax),%eax
	lea     '''+("__%s_TABLE_1"%modname)+'''(%eax),%eax
	push    %eax
	push    $'''+("__auto_iat_%s_handle"%modname)+'''
	push    $'''+("__auto_iat_%s_modname"%modname)+'''
	call    _AutoIatSpliceProcAddress@16
	xchg    %eax,(%esp)
	pop     (%eax)
	jmp     *(%eax)

.global '''+("__%s_xcall_fullinit"%modname)+'''
'''+("__%s_xcall_fullinit"%modname)+''':
	mov     $0, %eax
	push    %eax
_Loop:
	lea     '''+("__%s_TABLE"%modname)+'''(%eax),%eax
	cmpl    $'''+("__%s_xcall_resolver"%modname)+''',(%eax)
	jne     _Loop_Continue
	push    %eax
	mov     4(%esp),%eax
	push    %eax
	lea     '''+("__%s_TABLE_1"%modname)+'''(%eax),%eax
	mov     (%eax),%eax
	lea     '''+("__%s_TABLE_1"%modname)+'''(%eax),%eax
	push    %eax
	push    $'''+("__auto_iat_%s_handle"%modname)+'''
	push    $'''+("__auto_iat_%s_modname"%modname)+'''
	call    _AutoIatSpliceProcAddress@16
	xchg    %eax,(%esp)
	popl    (%eax)
_Loop_Continue:
	add     $4, (%esp)
	mov     (%esp),%eax
	cmpl    $('''+("__%s_TABLE_END"%modname)+'''-'''+("__%s_TABLE"%modname)+'''),%eax
	jl      _Loop
	pop     %eax
	ret

.data
/*.global '''+("__auto_iat_%s_handle"%modname)+'''*/
'''+("__auto_iat_%s_handle"%modname)+''':
.long 0
/*.global '''+("__auto_iat_%s_modname"%modname)+'''*/
'''+("__auto_iat_%s_modname"%modname)+''':
.asciz "'''+modname+'''"
''')

	def xor_text(x,sid):
		def next_rund(sid):
			n = 0
			while True:
				sid = 1664525*sid + 1013904223;
				yield int(sid) >> (n % 8)
				n +=1
		rund = next_rund(sid)
		r = ''.join(['\\x%02x' % ((ord(i) ^ rund.next())&0x0ff) for i in x])
		return r;

	fo_S.write(".global __%s_TABLE\n"%modname)
	fo_S.write("__%s_TABLE:\n"%modname)
	for i in xrange(fn_no):
		fo_S.write(".long __%s_xcall_resolver\n"%modname)
	fo_S.write(".global __%s_TABLE_END\n"%modname)
	fo_S.write("__%s_TABLE_END:\n"%modname)

	for i in xrange(fn_no):
		s = fn_names[i]
		j = xor_text(s+'\0',0xfabecaffe+(i*4))
		#j = s+'\0'
		fo_S.write('__%s_STRING_xx%d: .ascii "%s"\n' % (modname,i,j) )
	fo_S.write("__%s_TABLE_1:\n"%modname)
	for i in xrange(fn_no):
		fo_S.write('.long __%s_STRING_xx%d-__%s_TABLE_1\n' % (modname,i,modname) )

	for i,j in imp_names.items():
		fo_S.write(".text\n")
		fo_S.write(".global _%s\n"%j)
		fo_S.write("_%s:\n"%j)
		fo_S.write("__%s_xcall_resolver_%d:\n"%(modname,i))
		fo_S.write("    mov     $%d,%%eax\n"%i)
		fo_S.write("    jmp     __%s_xcall_resolver_X\n"%modname)
		fo_S.write(".data\n")
		fo_S.write(".global __imp__%s\n" % j)
		fo_S.write("__imp__%s:\n"%j)
		fo_S.write((".global __imp__%s\n" % j).replace('@','__'))
		fo_S.write(("__imp__%s:\n"%j).replace('@','__'))
		fo_S.write(".long __%s_xcall_resolver_%d\n"%(modname,i))

	for i,j in xcd_names.items():
		fo_S.write(".global _%s\n"%j)
		fo_S.write("_%s:\n"%j)
		fo_S.write("__%s_xcall_resolver_%d:\n"%(modname,i))
		fo_S.write("    mov     $%d,%%eax\n"%i)
		fo_S.write("    jmp     __%s_xcall_resolver_X\n"%modname)
		fo_S.write(".global __imp__%s\n" % j)
		fo_S.write("__imp__%s:\n"%j)
		fo_S.write(".long __%s_xcall_resolver_%d\n"%(modname,i))

Main(sys.argv)
